#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
资源文件打包脚本
将图片和模型文件复制到指定目录进行打包
"""

import os
import shutil
import glob
from pathlib import Path
import zipfile
import datetime

def create_package():
    """创建资源文件包"""
    
    # 创建打包目录
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    package_dir = f"assets_package_{timestamp}"
    os.makedirs(package_dir, exist_ok=True)
    
    print(f"📦 创建资源包目录: {package_dir}")
    
    # 1. 复制模型文件
    model_files = [
        "hrnet_w48_proto_lr1x_hrnet_proto_80k_latest.pth"
    ]
    
    models_dir = os.path.join(package_dir, "models")
    os.makedirs(models_dir, exist_ok=True)
    
    print("🔧 复制模型文件...")
    for model_file in model_files:
        if os.path.exists(model_file):
            shutil.copy2(model_file, models_dir)
            print(f"  ✅ {model_file}")
        else:
            print(f"  ❌ 未找到: {model_file}")
    
    # 2. 复制项目根目录的图片文件
    print("🖼️  复制项目根目录图片文件...")
    root_images = [
        "potsdam_label_visualization.png",
        "potsdam_sample_visualization.png",
        "street_test.jpg",
        "street_test2.jpg"
    ]
    
    images_dir = os.path.join(package_dir, "images")
    os.makedirs(images_dir, exist_ok=True)
    
    for img_file in root_images:
        if os.path.exists(img_file):
            shutil.copy2(img_file, images_dir)
            print(f"  ✅ {img_file}")
        else:
            print(f"  ❌ 未找到: {img_file}")
    
    # 3. 复制框架图片
    print("🖼️  复制框架图片...")
    framework_dir = os.path.join(package_dir, "images", "framework")
    os.makedirs(framework_dir, exist_ok=True)
    
    framework_files = glob.glob("figures/*.png")
    for fw_file in framework_files:
        shutil.copy2(fw_file, framework_dir)
        print(f"  ✅ {fw_file}")
    
    # 4. 复制部分处理后的数据样本（仅复制少量作为示例）
    print("📊 复制数据样本...")
    data_samples_dir = os.path.join(package_dir, "data_samples")
    os.makedirs(data_samples_dir, exist_ok=True)
    
    # 复制训练数据样本（每个类别选择几个）
    train_images_dir = os.path.join(data_samples_dir, "train_images")
    train_labels_dir = os.path.join(data_samples_dir, "train_labels")
    val_images_dir = os.path.join(data_samples_dir, "val_images")
    val_labels_dir = os.path.join(data_samples_dir, "val_labels")
    
    os.makedirs(train_images_dir, exist_ok=True)
    os.makedirs(train_labels_dir, exist_ok=True)
    os.makedirs(val_images_dir, exist_ok=True)
    os.makedirs(val_labels_dir, exist_ok=True)
    
    # 复制训练数据样本（每个目录最多5个文件）
    sample_count = 0
    max_samples = 5
    
    for root, dirs, files in os.walk("data/potsdam/processed/train/images"):
        for file in files[:max_samples]:
            if file.endswith('.png'):
                src = os.path.join(root, file)
                dst = os.path.join(train_images_dir, file)
                shutil.copy2(src, dst)
                sample_count += 1
                if sample_count >= max_samples:
                    break
        if sample_count >= max_samples:
            break
    
    sample_count = 0
    for root, dirs, files in os.walk("data/potsdam/processed/train/labels"):
        for file in files[:max_samples]:
            if file.endswith('.png'):
                src = os.path.join(root, file)
                dst = os.path.join(train_labels_dir, file)
                shutil.copy2(src, dst)
                sample_count += 1
                if sample_count >= max_samples:
                    break
        if sample_count >= max_samples:
            break
    
    # 复制验证数据样本
    sample_count = 0
    for root, dirs, files in os.walk("data/potsdam/processed/val/images"):
        for file in files[:max_samples]:
            if file.endswith('.png'):
                src = os.path.join(root, file)
                dst = os.path.join(val_images_dir, file)
                shutil.copy2(src, dst)
                sample_count += 1
                if sample_count >= max_samples:
                    break
        if sample_count >= max_samples:
            break
    
    sample_count = 0
    for root, dirs, files in os.walk("data/potsdam/processed/val/labels"):
        for file in files[:max_samples]:
            if file.endswith('.png'):
                src = os.path.join(root, file)
                dst = os.path.join(val_labels_dir, file)
                shutil.copy2(src, dst)
                sample_count += 1
                if sample_count >= max_samples:
                    break
        if sample_count >= max_samples:
            break
    
    print(f"  ✅ 复制了 {max_samples} 个训练和验证数据样本")
    
    # 5. 创建说明文件
    readme_content = f"""# ProtoSeg Potsdam 资源包

## 创建时间
{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

## 包含内容

### 模型文件 (models/)
- hrnet_w48_proto_lr1x_hrnet_proto_80k_latest.pth: 预训练模型

### 图片文件 (images/)
- potsdam_label_visualization.png: 标签可视化
- potsdam_sample_visualization.png: 样本可视化
- street_test.jpg, street_test2.jpg: 测试图片
- framework/: 框架图片

### 数据样本 (data_samples/)
- train_images/: 训练图像样本
- train_labels/: 训练标签样本
- val_images/: 验证图像样本
- val_labels/: 验证标签样本

## 使用说明
1. 模型文件可用于模型加载和推理
2. 图片文件可用于测试和可视化
3. 数据样本可用于理解数据格式和结构

## 注意事项
- 此包仅包含部分数据样本，完整数据集请从原始数据源获取
- 模型文件较大，请确保有足够的存储空间
"""
    
    readme_path = os.path.join(package_dir, "README.md")
    with open(readme_path, 'w', encoding='utf-8') as f:
        f.write(readme_content)
    
    print("📝 创建说明文件: README.md")
    
    # 6. 创建压缩包
    print("🗜️  创建压缩包...")
    zip_filename = f"{package_dir}.zip"
    
    with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(package_dir):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, package_dir)
                zipf.write(file_path, arcname)
    
    print(f"✅ 压缩包创建完成: {zip_filename}")
    
    # 7. 统计信息
    total_size = 0
    file_count = 0
    
    for root, dirs, files in os.walk(package_dir):
        for file in files:
            file_path = os.path.join(root, file)
            total_size += os.path.getsize(file_path)
            file_count += 1
    
    zip_size = os.path.getsize(zip_filename)
    
    print("\n📊 打包统计:")
    print(f"  文件数量: {file_count}")
    print(f"  目录大小: {total_size / (1024*1024):.2f} MB")
    print(f"  压缩包大小: {zip_size / (1024*1024):.2f} MB")
    print(f"  压缩率: {(1 - zip_size/total_size)*100:.1f}%")
    
    print(f"\n🎉 资源包创建完成!")
    print(f"📁 目录: {package_dir}")
    print(f"📦 压缩包: {zip_filename}")
    
    return package_dir, zip_filename

if __name__ == "__main__":
    create_package()
